<!DOCTYPE html>
<html lang="en">
	<head>
		<title>three.js xr - dragging with custom depth shader</title>
		<meta charset="utf-8">
		<meta name="viewport" content="width=device-width, initial-scale=1.0, user-scalable=no">
		<link type="text/css" rel="stylesheet" href="main.css">
	</head>
	<body>

		<div id="info">
			<a href="https://threejs.org" target="_blank" rel="noopener">three.js</a> xr - dragging
		</div>

		<script type="importmap">
			{
				"imports": {
					"three": "../build/three.module.js",
					"three/addons/": "./jsm/"
				}
			}
		</script>

		<script type="module">

			import * as THREE from 'three';
			import { OrbitControls } from 'three/addons/controls/OrbitControls.js';
			import { XRButton } from 'three/addons/webxr/XRButton.js';
			import { XRControllerModelFactory } from 'three/addons/webxr/XRControllerModelFactory.js';

			let container;
			let camera, scene, renderer;
			let controller1, controller2;
			let controllerGrip1, controllerGrip2;
			let isDepthSupplied = false;

			let raycaster;

			const intersected = [];
			const tempMatrix = new THREE.Matrix4();

			let controls, group;

			init();
			animate();

			function init() {

				container = document.createElement( 'div' );
				document.body.appendChild( container );

				scene = new THREE.Scene();
				scene.background = new THREE.Color( 0x808080 );

				camera = new THREE.PerspectiveCamera( 50, window.innerWidth / window.innerHeight, 0.1, 10 );
				camera.position.set( 0, 1.6, 3 );

				controls = new OrbitControls( camera, container );
				controls.target.set( 0, 1.6, 0 );
				controls.update();

				const floorGeometry = new THREE.PlaneGeometry( 6, 6 );
				const floorMaterial = new THREE.ShadowMaterial( { opacity: 0.25, blending: THREE.CustomBlending, transparent: false } );
				const floor = new THREE.Mesh( floorGeometry, floorMaterial );
				floor.rotation.x = - Math.PI / 2;
				floor.receiveShadow = true;
				scene.add( floor );

				scene.add( new THREE.HemisphereLight( 0xbcbcbc, 0xa5a5a5, 3 ) );

				const light = new THREE.DirectionalLight( 0xffffff, 3 );
				light.position.set( 0, 6, 0 );
				light.castShadow = true;
				light.shadow.camera.top = 3;
				light.shadow.camera.bottom = - 3;
				light.shadow.camera.right = 3;
				light.shadow.camera.left = - 3;
				light.shadow.mapSize.set( 4096, 4096 );
				scene.add( light );

				group = new THREE.Group();
				scene.add( group );

				const geometries = [
					new THREE.BoxGeometry( 0.2, 0.2, 0.2 ),
					new THREE.ConeGeometry( 0.2, 0.2, 64 ),
					new THREE.CylinderGeometry( 0.2, 0.2, 0.2, 64 ),
					new THREE.IcosahedronGeometry( 0.2, 8 ),
					new THREE.TorusGeometry( 0.2, 0.04, 64, 32 )
				];

				for ( let i = 0; i < 50; i ++ ) {

					const geometry = geometries[ Math.floor( Math.random() * geometries.length ) ];
					const material = new THREE.ShaderMaterial( {

						vertexShader: /* glsl */`
							varying vec3 vNormal;
							varying vec2 vUv;
							void main() {
								vNormal = normalize(normalMatrix * normal);
								vUv = uv;
								gl_Position = projectionMatrix * modelViewMatrix * vec4(position, 1.0);
							}
						`,

						fragmentShader: /* glsl */`
							uniform vec3 diffuseColor;
							uniform float roughness;
							uniform float metalness;
							uniform float emissive;
							varying vec3 vNormal;
							varying vec2 vUv;
							uniform sampler2DArray depthColor;
							uniform float depthWidth;
							uniform float depthHeight;
							#define saturate( a ) clamp( a, 0.0, 1.0 )
							float Depth_GetCameraDepthInMeters(const sampler2DArray depthTexture,
								const vec2 depthUv, int arrayIndex) {
								return texture(depthColor, vec3(depthUv.x, depthUv.y, arrayIndex)).r;
							}
							float Depth_GetOcclusion(const sampler2DArray depthTexture, const vec2 depthUv, float assetDepthM, int arrayIndex) {
								float depthMm = Depth_GetCameraDepthInMeters(depthTexture, depthUv, arrayIndex);
								const float kDepthTolerancePerM = 0.001;
								return clamp(1.0 -
									0.5 * (depthMm - assetDepthM) /
										(kDepthTolerancePerM * assetDepthM) +
									0.5, 0.0, 1.0);
							}
							float Depth_GetBlurredOcclusionAroundUV(const sampler2DArray depthTexture, const vec2 uv, float assetDepthM, int arrayIndex) {
								// Kernel used:
								// 0   4   7   4   0
								// 4   16  26  16  4
								// 7   26  41  26  7
								// 4   16  26  16  4
								// 0   4   7   4   0
								const float kKernelTotalWeights = 269.0;
								float sum = 0.0;
								const float kOcclusionBlurAmount = 0.0005;
								vec2 blurriness =
								vec2(kOcclusionBlurAmount, kOcclusionBlurAmount /** u_DepthAspectRatio*/);
								float current = 0.0;
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(-1.0, -2.0) * blurriness, assetDepthM, arrayIndex);
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(+1.0, -2.0) * blurriness, assetDepthM, arrayIndex);
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(-1.0, +2.0) * blurriness, assetDepthM, arrayIndex);
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(+1.0, +2.0) * blurriness, assetDepthM, arrayIndex);
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(-2.0, +1.0) * blurriness, assetDepthM, arrayIndex);
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(+2.0, +1.0) * blurriness, assetDepthM, arrayIndex);
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(-2.0, -1.0) * blurriness, assetDepthM, arrayIndex);
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(+2.0, -1.0) * blurriness, assetDepthM, arrayIndex);
								sum += current * 4.0;
								current = 0.0;
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(-2.0, -0.0) * blurriness, assetDepthM, arrayIndex);
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(+2.0, +0.0) * blurriness, assetDepthM, arrayIndex);
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(+0.0, +2.0) * blurriness, assetDepthM, arrayIndex);
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(-0.0, -2.0) * blurriness, assetDepthM, arrayIndex);
								sum += current * 7.0;
								current = 0.0;
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(-1.0, -1.0) * blurriness, assetDepthM, arrayIndex);
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(+1.0, -1.0) * blurriness, assetDepthM, arrayIndex);
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(-1.0, +1.0) * blurriness, assetDepthM, arrayIndex);
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(+1.0, +1.0) * blurriness, assetDepthM, arrayIndex);
								sum += current * 16.0;
								current = 0.0;
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(+0.0, +1.0) * blurriness, assetDepthM, arrayIndex);
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(-0.0, -1.0) * blurriness, assetDepthM, arrayIndex);
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(-1.0, -0.0) * blurriness, assetDepthM, arrayIndex);
								current += Depth_GetOcclusion(
								depthTexture, uv + vec2(+1.0, +0.0) * blurriness, assetDepthM, arrayIndex);
								sum += current * 26.0;
								sum += Depth_GetOcclusion(depthTexture, uv, assetDepthM, arrayIndex) * 41.0;
								return sum / kKernelTotalWeights;
							}
							void main() {
								vec3 normal = normalize(vNormal);
								vec3 diffuse = diffuseColor;
								float specularIntensity = pow(max(dot(normal, normalize(vec3(0, 0, 1))), 0.0), 64.0);
								vec3 specular = vec3(specularIntensity) * mix(vec3(0.04), diffuse, metalness);
								gl_FragColor = vec4(diffuse * (1.0 - specular) + specular, 1.0) * (1.0 + emissive);

								if (depthWidth > 0.0) {
									int arrayIndex = 0;
									vec2 depthUv;
									if (gl_FragCoord.x>=depthWidth) {
										arrayIndex = 1;
										depthUv = vec2((gl_FragCoord.x-depthWidth)/depthWidth, gl_FragCoord.y/depthHeight);
									} else {
										depthUv = vec2(gl_FragCoord.x/depthWidth, gl_FragCoord.y/depthHeight);
									}
									float assetDepthM = gl_FragCoord.z;

									float occlusion = Depth_GetBlurredOcclusionAroundUV(depthColor, depthUv, assetDepthM, arrayIndex);
									float depthMm = Depth_GetCameraDepthInMeters(depthColor, depthUv, arrayIndex);

									float absDistance = abs(assetDepthM - depthMm);
									float v = 0.0025;
									absDistance = saturate(v - absDistance) / v;

									gl_FragColor.rgb += vec3(absDistance * 2.0, absDistance * 2.0, absDistance * 12.0);
									gl_FragColor = mix(gl_FragColor, vec4(0.0, 0.0, 0.0, 0.0), occlusion * 0.7);
								}
							}
							`,

						uniforms: {

							diffuseColor: { value: new THREE.Color( Math.random() * 0xffffff ) },
							roughness: { value: 0.7 },
							metalness: { value: 0.0 },
							emissive: { value: 0.0 },
							depthWidth: { value: 0.0 },
							depthHeight: { value: 0.0 },
							depthColor: { value: new THREE.Texture() }

						}

					} );

					const object = new THREE.Mesh( geometry, material );

					object.position.x = Math.random() * 4 - 2;
					object.position.y = Math.random() * 2;
					object.position.z = Math.random() * 4 - 2;

					object.rotation.x = Math.random() * 2 * Math.PI;
					object.rotation.y = Math.random() * 2 * Math.PI;
					object.rotation.z = Math.random() * 2 * Math.PI;

					object.scale.setScalar( Math.random() + 0.5 );

					group.add( object );

				}

				//

				renderer = new THREE.WebGLRenderer( { antialias: true } );
				renderer.setPixelRatio( window.devicePixelRatio );
				renderer.setSize( window.innerWidth, window.innerHeight );
				renderer.shadowMap.enabled = true;
				renderer.xr.enabled = true;
				container.appendChild( renderer.domElement );

				document.body.appendChild( XRButton.createButton( renderer, {
					'optionalFeatures': [ 'depth-sensing' ],
					'depthSensing': { 'usagePreference': [ 'gpu-optimized' ], 'dataFormatPreference': [] }
				} ) );

				// controllers

				controller1 = renderer.xr.getController( 0 );
				controller1.addEventListener( 'selectstart', onSelectStart );
				controller1.addEventListener( 'selectend', onSelectEnd );
				scene.add( controller1 );

				controller2 = renderer.xr.getController( 1 );
				controller2.addEventListener( 'selectstart', onSelectStart );
				controller2.addEventListener( 'selectend', onSelectEnd );
				scene.add( controller2 );

				const controllerModelFactory = new XRControllerModelFactory();

				controllerGrip1 = renderer.xr.getControllerGrip( 0 );
				controllerGrip1.add( controllerModelFactory.createControllerModel( controllerGrip1 ) );
				scene.add( controllerGrip1 );

				controllerGrip2 = renderer.xr.getControllerGrip( 1 );
				controllerGrip2.add( controllerModelFactory.createControllerModel( controllerGrip2 ) );
				scene.add( controllerGrip2 );

				//

				const geometry = new THREE.BufferGeometry().setFromPoints( [ new THREE.Vector3( 0, 0, 0 ), new THREE.Vector3( 0, 0, - 1 ) ] );

				const line = new THREE.Line( geometry );
				line.name = 'line';
				line.scale.z = 5;

				controller1.add( line.clone() );
				controller2.add( line.clone() );

				raycaster = new THREE.Raycaster();

				//

				window.addEventListener( 'resize', onWindowResize );

			}

			function onWindowResize() {

				camera.aspect = window.innerWidth / window.innerHeight;
				camera.updateProjectionMatrix();

				renderer.setSize( window.innerWidth, window.innerHeight );

			}

			function onSelectStart( event ) {

				const controller = event.target;

				const intersections = getIntersections( controller );

				if ( intersections.length > 0 ) {

					const intersection = intersections[ 0 ];

					const object = intersection.object;
					object.material.uniforms.emissive.value = 1;
					controller.attach( object );

					controller.userData.selected = object;

				}

				controller.userData.targetRayMode = event.data.targetRayMode;

			}

			function onSelectEnd( event ) {

				const controller = event.target;

				if ( controller.userData.selected !== undefined ) {

					const object = controller.userData.selected;
					object.material.uniforms.emissive.value = 0;
					group.attach( object );

					controller.userData.selected = undefined;

				}

			}

			function getIntersections( controller ) {

				controller.updateMatrixWorld();

				tempMatrix.identity().extractRotation( controller.matrixWorld );

				raycaster.ray.origin.setFromMatrixPosition( controller.matrixWorld );
				raycaster.ray.direction.set( 0, 0, - 1 ).applyMatrix4( tempMatrix );

				return raycaster.intersectObjects( group.children, false );

			}

			function intersectObjects( controller ) {

				// Do not highlight in mobile-ar

				if ( controller.userData.targetRayMode === 'screen' ) return;

				// Do not highlight when already selected

				if ( controller.userData.selected !== undefined ) return;

				const line = controller.getObjectByName( 'line' );
				const intersections = getIntersections( controller );

				if ( intersections.length > 0 ) {

					const intersection = intersections[ 0 ];

					const object = intersection.object;
					object.material.uniforms.emissive.value = 1;
					intersected.push( object );

					line.scale.z = intersection.distance;

				} else {

					line.scale.z = 5;

				}

			}

			function cleanIntersected() {

				while ( intersected.length ) {

					const object = intersected.pop();
					object.material.uniforms.emissive.value = 0;

				}

			}

			//

			function animate() {

				renderer.setAnimationLoop( render );

			}

			function render() {

				if ( renderer.xr.hasDepthSensing() && ! isDepthSupplied ) {

					group.children.forEach( child => {

						child.material.uniforms.depthColor.value = renderer.xr.getDepthTexture();
						child.material.uniforms.depthWidth.value = 1680;
						child.material.uniforms.depthHeight.value = 1760;

						isDepthSupplied = true;

					} );

				} else if ( ! renderer.xr.hasDepthSensing() && isDepthSupplied ) {

					group.children.forEach( child => {

						child.material.uniforms.depthWidth.value = 0;
						child.material.uniforms.depthHeight.value = 0;

						isDepthSupplied = false;

					} );

				}

				cleanIntersected();

				intersectObjects( controller1 );
				intersectObjects( controller2 );

				renderer.render( scene, camera );

			}

		</script>
	</body>
</html>
